#include "gpu/esdf_map.h"

#include <cstdlib>
#include <fstream>  // 文件流操作
#include <iostream>

#include "common/config.h"
#include "gpu/kernel.h"

namespace {
inline void checkCudaError(cudaError_t result, const char* message) {
    if (result != cudaSuccess) {
        std::cerr << "CUDA error: " << message << " - " << cudaGetErrorString(result) << std::endl;
        exit(EXIT_FAILURE);
    }
}

template <class T>
void delete_arr(T& ptr) {
    if (ptr != nullptr) {
        delete[] ptr;
        ptr = nullptr;
    }
}
}  // namespace

ESDFMap::ESDFMap(int map_row_size, int map_col_size, float resolution) : map_row_size_(map_row_size), map_col_size_(map_col_size), resolution_(resolution) {
    int size = map_row_size_ * map_col_size_;
    occupancy_buffer_inflate_ = new uint8_t[size];
    tmp_buffer_ = new float[size];
    distance_buffer_ = new float[size];
    obs_index_data_ = new int[kMaxObstacleGridNums];
    obs_start_index_data_ = new int[map_row_size_ + 1];

    checkCudaError(cudaStreamCreate(&stream_), "cudaStreamCreate");

    size_t used_gpu_memory = 0;
    cudaMalloc(&occupancy_buffer_inflate_gpu_, size * sizeof(uint8_t));
    used_gpu_memory += size * sizeof(uint8_t);
    cudaMalloc(&tmp_buffer_gpu_, size * sizeof(float));
    used_gpu_memory += size * sizeof(float);
    cudaMalloc(&distance_buffer_gpu_, size * sizeof(float));
    used_gpu_memory += size * sizeof(float);
    cudaMalloc(&obs_index_data_gpu_, kMaxObstacleGridNums * sizeof(int));
    used_gpu_memory += kMaxObstacleGridNums * sizeof(int);
    cudaMalloc(&obs_start_index_data_gpu_, (map_row_size_ + 1) * sizeof(int));
    used_gpu_memory += (map_row_size_ + 1) * sizeof(int);
    std::cout << "used_gpu_memory: " << static_cast<float>(used_gpu_memory) / 1024 / 1024 << "MB" << std::endl;

    distance_buffer_gpu_res_ = new float[size];
    reset();
}

ESDFMap::~ESDFMap() {
    delete_arr(occupancy_buffer_inflate_);
    delete_arr(tmp_buffer_);
    delete_arr(distance_buffer_);
    delete_arr(obs_index_data_);
    delete_arr(obs_start_index_data_);
    delete_arr(distance_buffer_gpu_res_);

    cudaFree(occupancy_buffer_inflate_gpu_);
    cudaFree(tmp_buffer_gpu_);
    cudaFree(distance_buffer_gpu_);
    cudaFree(obs_index_data_gpu_);
    cudaFree(obs_start_index_data_gpu_);
}

void ESDFMap::reset() {
    int size = map_row_size_ * map_col_size_;
    std::memset(occupancy_buffer_inflate_, static_cast<uint8_t>(OccupancyType::NONE), size * sizeof(uint8_t));
    std::memset(tmp_buffer_, 0, size * sizeof(float));
    std::memset(distance_buffer_, 10000.0, size * sizeof(float));  // 填一个比较大的值, 方便debug
}

void ESDFMap::init_occupancy_boundary(const std::vector<Polygon>& polygon_vec) {
    for (const auto& polygon : polygon_vec) {
        const std::vector<Point2d>& corner_points = polygon.corner_points();
        for (int i = 0; i < corner_points.size(); i++) {
            const auto& p0 = corner_points[i];
            const auto& p1 = corner_points[(i + 1) % corner_points.size()];
            bresenham_line(p0.x, p0.y, p1.x, p1.y);
        }
    }
}

void ESDFMap::init_occupancy_inside(const std::vector<Polygon>& polygon_vec) {
    for (const auto& polygon : polygon_vec) {
        const std::vector<Point2d>& corner_points = polygon.corner_points();
        bfs_set_occupancy(polygon.internal_point(), OccupancyType::INSIDE);
    }
}

void ESDFMap::calc_esdf_cpu() {
    for (int x = 0; x < map_row_size_; x++) {  // 行扫描
        fill_esdf(
            [&](int y) {
                return occupancy_buffer_inflate_[to_address(x, y)] == static_cast<uint8_t>(OccupancyType::BOUNDARY) ? 0.0 : std::numeric_limits<float>::max();
            },
            [&](int y, float val) { tmp_buffer_[to_address(x, y)] = val; }, 0, map_col_size_ - 1, map_col_size_);
    }

    for (int y = 0; y < map_col_size_; y++) {  // 列扫描
        fill_esdf([&](int x) { return tmp_buffer_[to_address(x, y)]; },
                  [&](int x, float val) { distance_buffer_[to_address(x, y)] = resolution_ * std::sqrt(val); }, 0, map_row_size_ - 1, map_row_size_);
    }
}

void ESDFMap::calc_esdf_gpu() {
    int obs_index = 0;
    for (int x = 0; x < map_row_size_; x++) {
        obs_start_index_data_[x] = obs_index;
        // std::cout << "x: " << x<< ", obs_start_index: " << obs_index << std::endl;
        for (int y = 0; y < map_col_size_; y++) {
            if (get_occupancy(x, y) == static_cast<uint8_t>(OccupancyType::BOUNDARY)) {
                assert(obs_index < kMaxObstacleGridNums);
                obs_index_data_[obs_index] = y;
                // std::cout << "(" << obs_index << ", " << y <<"), ";
                obs_index++;
            }
        }
        // std::cout << std::endl;
    }
    obs_start_index_data_[map_row_size_] = obs_index;

    // copy一些数据到gpu上
    int size = map_row_size_ * map_col_size_;
    checkCudaError(cudaMemcpyAsync(occupancy_buffer_inflate_gpu_, occupancy_buffer_inflate_, size * sizeof(uint8_t), cudaMemcpyHostToDevice, stream_),
                   "cudaMemcpyAsync (HostToDevice)");
    checkCudaError(cudaMemcpyAsync(tmp_buffer_gpu_, tmp_buffer_, size * sizeof(float), cudaMemcpyHostToDevice, stream_), 
                   "cudaMemcpyAsync (HostToDevice)");
    checkCudaError(cudaMemcpyAsync(obs_index_data_gpu_, obs_index_data_, kMaxObstacleGridNums * sizeof(int), cudaMemcpyHostToDevice, stream_),
                   "cudaMemcpyAsync (HostToDevice)");
    checkCudaError(cudaMemcpyAsync(obs_start_index_data_gpu_, obs_start_index_data_, (map_row_size_ + 1) * sizeof(int), cudaMemcpyHostToDevice, stream_),
                   "cudaMemcpyAsync (HostToDevice)");

    // 在gpu上进行计算
    calc_esdf_map_launcher(&stream_, map_row_size_, map_col_size_, resolution_, occupancy_buffer_inflate_gpu_, tmp_buffer_gpu_, distance_buffer_gpu_,
                           obs_index_data_gpu_, obs_start_index_data_gpu_);

    // 将计算的结果copy回cpu
    checkCudaError(cudaMemcpyAsync(distance_buffer_gpu_res_, distance_buffer_gpu_, size * sizeof(float), cudaMemcpyDeviceToHost, stream_),
                   "cudaMemcpyAsync (DeviceToHost)");

    checkCudaError(cudaStreamSynchronize(stream_), "cudaStreamSynchronize");
}

void ESDFMap::debug_print_occupancy() const {
    std::ofstream out_file("../res/occupancy_data.txt");
    out_file << map_row_size_ << "\n";
    out_file << map_col_size_ << "\n";

    out_file << "BOUNDARY: ";
    for (int row = 0; row < map_row_size_; row++) {
        for (int col = 0; col < map_col_size_; col++) {
            if (get_occupancy(row, col) == static_cast<uint8_t>(OccupancyType::BOUNDARY)) {
                out_file << "(" << row << ", " << col << "), ";
            }
        }
    }
    out_file << "\n";
    out_file << "INSIDE: ";
    for (int row = 0; row < map_row_size_; row++) {
        for (int col = 0; col < map_col_size_; col++) {
            if (get_occupancy(row, col) == static_cast<uint8_t>(OccupancyType::INSIDE)) {
                out_file << "(" << row << ", " << col << "), ";
            }
        }
    }

    out_file.close();
}

void ESDFMap::debug_print_esdf_cpu() const {
    std::ofstream out_file2("../res/esdf_cpu.txt");
    out_file2 << map_row_size_ << "\n";
    out_file2 << map_col_size_ << "\n";
    for (int row = 0; row < map_row_size_; row++) {
        for (int col = 0; col < map_col_size_; col++) {
            float distance = get_distance_cpu(row, col);
            out_file2 << "(" << row << ", " << col << ", " << distance << "), ";
        }
    }
    out_file2.close();
}

void ESDFMap::debug_print_esdf_gpu() const {
    std::ofstream out_file3("../res/esdf_gpu.txt");
    out_file3 << map_row_size_ << "\n";
    out_file3 << map_col_size_ << "\n";
    for (int row = 0; row < map_row_size_; row++) {
        for (int col = 0; col < map_col_size_; col++) {
            float distance = get_distance_gpu(row, col);
            out_file3 << "(" << row << ", " << col << ", " << distance << "), ";
        }
    }
    out_file3.close();
}

// 在计算机图形学中，将连续的直线离散化为像素点的过程称为直线光栅化。
// 最常用的算法是Bresenham 算法，它通过整数运算高效地确定直线经过的像素点。
void ESDFMap::bresenham_line(float x0, float y0, float x1, float y1) {
    // 转换为网格索引（整数）
    int ix0 = float_to_grid(x0);
    int iy0 = float_to_grid(y0);
    int ix1 = float_to_grid(x1);
    int iy1 = float_to_grid(y1);

    int dx = std::abs(ix1 - ix0);
    int dy = std::abs(iy1 - iy0);
    int step_x = (ix0 < ix1) ? 1 : -1;
    int step_y = (iy0 < iy1) ? 1 : -1;
    int err = dx - dy;

    while (true) {
        if (!is_index_valid(ix0, iy0)) {
            continue;
        }

        occupancy_buffer_inflate_[to_address(ix0, iy0)] = static_cast<uint8_t>(OccupancyType::BOUNDARY);

        if (ix0 == ix1 && iy0 == iy1)
            break;

        int e2 = 2 * err;
        if (e2 > -dy) {
            err -= dy;
            ix0 += step_x;
        }
        if (e2 < dx) {
            err += dx;
            iy0 += step_y;
        }
    }
}

void ESDFMap::bfs_set_occupancy(const Point2d& internal_point, OccupancyType flag) {
    assert(flag != OccupancyType::NONE);
    const static std::vector<std::pair<int, int>> directions = {{-1, 0}, {1, 0}, {0, -1}, {0, 1}};

    int start_x = float_to_grid(internal_point.x);
    int start_y = float_to_grid(internal_point.y);
    if (!is_index_valid(start_x, start_y)) {
        return;
    }

    std::queue<std::pair<int, int>> queue;  // BFS 队列
    queue.push({start_x, start_y});
    occupancy_buffer_inflate_[to_address(start_x, start_y)] = static_cast<uint8_t>(flag);
    // std::cout << "start_x/start_y: " << start_x << ", " << start_y << std::endl;

    while (!queue.empty()) {
        auto [x, y] = queue.front();
        queue.pop();

        // 遍历四个方向
        for (const auto& [dx, dy] : directions) {
            int newX = x + dx;
            int newY = y + dy;

            if (is_index_valid(newX, newY) && occupancy_buffer_inflate_[to_address(newX, newY)] == static_cast<uint8_t>(OccupancyType::NONE)) {
                queue.push({newX, newY});
                occupancy_buffer_inflate_[to_address(newX, newY)] = static_cast<uint8_t>(flag);
            }
        }
    }
}